import torch
import os

# 假设 ABINet 模型定义在 abinet.models.abinet 下
from modules.model_abinet import ABINetModel as ABINet
# from abinet.models.abinet import ABINet
# 假设有自定义的数据集加载器
from abinet.datasets import get_dataloader

def finetune_abinet(
    pretrained_path,
    save_path,
    train_data_dir,
    val_data_dir=None,
    epochs=5,
    lr=1e-4,
    batch_size=32,
    device='cpu'
):
    # 1. 加载模型结构
    model = ABINet()
    model = model.to(device)

    # 2. 加载预训练权重
    checkpoint = torch.load(pretrained_path, map_location=device)
    model.load_state_dict(checkpoint['state_dict'] if 'state_dict' in checkpoint else checkpoint)

    # 3. 数据加载
    train_loader = get_dataloader(train_data_dir, batch_size=batch_size, shuffle=True)
    if val_data_dir:
        val_loader = get_dataloader(val_data_dir, batch_size=batch_size, shuffle=False)
    else:
        val_loader = None

    # 4. 优化器和损失
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = torch.nn.CTCLoss()  # 视具体任务而定

    # 5. 微调训练
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for images, labels, label_lengths in train_loader:
            images = images.to(device)
            labels = labels.to(device)
            # 前向
            outputs = model(images)
            # 计算损失
            loss = criterion(outputs, labels, ...)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}")

        # 可选：验证集评估
        # ...

    # 6. 保存新模型
    torch.save(model.state_dict(), save_path)
    print(f"Finetuned model saved to {save_path}")

if __name__ == '__main__':
    pretrained_path = r'D:\songlin\Python\26_检修后续实验\ABINet\workdir\train-abinet\best-train-abinet.pth'
    save_path = r'D:\songlin\Python\26_检修后续实验\ABINet\workdir\train-abinet\finetuned-abinet.pth'
    train_data_dir = r'D:\songlin\Python\26_检修后续实验\ABINet\data\training\MJ\temp_test\train'
    # val_data_dir = r'你的新验证集路径'  # 可选
    finetune_abinet(pretrained_path, save_path, train_data_dir)
